#pragma once

#include <stdlib.h>
#include <stdio.h>
#include "structures.h"

typedef struct _BitStream {
	unsigned char* base;
	unsigned char* ptr;
	unsigned char offset;
	size_t capacity;
} BitStream;

void err(char* err)
{
	printf("[ERROR MESSAGE]: %s", err);
	fflush(stdout);
	exit(-1);
}

void bsshift1(BitStream* bs)
{
	/* shift the offset of bitstream bs */
	bs->offset++;
	if (bs->offset % (sizeof(unsigned char) * 8) == 0) {
		bs->ptr++;
		bs->offset = 0;
	}
}

BitStream* bsnew(unsigned char* buffer, size_t cap)
{
	BitStream* bs = (BitStream*)malloc(sizeof(BitStream));
	bs->base = buffer;
	bs->ptr = buffer;
	bs->capacity = cap;
	bs->offset = 0;
	return bs;
}

unsigned char bsgetb(BitStream* bs)
{
	/* get the LSB of the bitstream bs */
	unsigned char bit = (*(bs->ptr) & (1 << bs->offset)) >> bs->offset;
	bsshift1(bs);
	return bit;
}

void bswriteb(BitStream* bs, unsigned char b)
{
	/* write the least significant bit of b into the bitstrem bs (LSB) */
	*(bs->ptr) &= ((unsigned char)0xff >> (8 - bs->offset));
	if (b & 1)
		* (bs->ptr) |= (1 << bs->offset);
	bsshift1(bs);
}

void bswritec(BitStream* bs, Freqperchar* freq)
{
	size_t cnt = freq->encodebits;
	for (; cnt; cnt--)
		bswriteb(bs, (unsigned char)(freq->code >> (cnt - 1)));
}

Freqperchar* bsgetc_huff(BitStream* bs, const TreeNode* tree)
{
	if ((size_t)(bs->ptr - bs->base) >= bs->capacity)
		err((char *)"BitStream reached end!\n");
	TreeNode* cptr = (TreeNode*)tree;
	while (1) {
		if (!(bsgetb(bs) & 1))
			cptr = cptr->left;
		else
			cptr = cptr->right;
		if (cptr->leaf && !cptr->left && !cptr->right)
			return cptr->leaf;
	}
}

void bsdebug(BitStream* bs)
{
	printf("bs now is %d %d %d %d %d %d %d %d", bs->base[0], bs->base[1], bs->base[2], bs->base[3],
		bs->base[4], bs->base[5], bs->base[6], bs->base[7]);
	printf("\n");
}